
import os
import numpy as np
from pathlib import Path
from scipy.io import loadmat,savemat

COMBINE_COEFF_RBGFACE = True
COMBINE_LAMDK_RBGFACE = True


if __name__ == "__main__":

    path_data = Path("/mnt/data/DATA/LRW/processing")


    #num_train_file = 3000
    #num_train_count = 0

    # for each word
    v_ids = sorted(path_data.glob("*"))
    print(f"Have {len(v_ids)} dir files, \n {v_ids[:3]}")

    for i, ipath_id in enumerate(v_ids):
        id_name = ipath_id.stem

        dir_video = path_data.joinpath(id_name)
        v_dir_videos = sorted(dir_video.glob("*/*/video/video_orignal.mp4"))

        print(f"\nid= {id_name}, Have {len(v_dir_videos)} video files, \n {v_dir_videos[:3]}")

        # for each video
        for j, jpath_videos in enumerate(v_dir_videos):

            jpath_data = jpath_videos.parent.parent
            #print(f"jpath_data = {jpath_data}")

            jpath_frames = jpath_data.joinpath("frames")
            v_frames = sorted(jpath_frames.glob("*.jpg"), key=lambda x: int(x.stem[-4:]))
            n_frames = len(v_frames)
            #print(f"contain {len(v_frames)} frames in {str(jpath_data)}")


            jpath_coeff = jpath_data.joinpath("reconstruction/coeff")
            jpath_coeff_save = jpath_data.joinpath("video/coeff_reconstruction.npy")   
            jpath_coeff_save.parent.mkdir(parents=True, exist_ok=True)

            n_coeff = len(list(jpath_coeff.glob("*.txt")))
            #print(f"contain {len(v_coeff)} reconstruction/tans in {str(jpath_data)}")

            if n_frames != n_coeff:
                print(f"\t ERROR && Continue: This video missing coeff {n_coeff}/{n_frames}, {str(jpath_coeff)}")
                continue
            # elif jpath_coeff_save.exists():
            #     continue
            else:
                coeff_all = np.zeros((n_frames, 70), np.float32)  # 80 identity + 64 expression + 80 texture + 3 rot + 27 light + 3 trans
                for k, iframe in enumerate(v_frames):
                    fid = int(iframe.stem[-4:])
                        
                    # load data
                    kpath_coeff = jpath_coeff.joinpath(f"frame-{fid:04d}.txt")
                    kk_coeff = np.loadtxt(str(kpath_coeff)).reshape(-1)

                    kk_coeff_exp = kk_coeff[ 80:144]     # (64, ) , expression-64      ==> (55, )
                    kk_coeff_rot = kk_coeff[224:227]     # (3, ),   rotation-3, 
                    kk_coeff_tra = kk_coeff[254:257]     # (3, ),   translation - 3

                    coeff_all[k, :] = np.r_[kk_coeff_exp, kk_coeff_rot, kk_coeff_tra]      #  (70, )  expression-55 + rotation-3 + trans-3

                    #print(f"coef shape={coeff_all.shape}, path = {kpath_coeff}")
                # save combine data
                np.save(str(jpath_coeff_save), coeff_all)

                if j % 100 == 99:
                    print(f"processing {j+1}/{len(v_dir_videos)}, shape={coeff_all.shape} --> {str(jpath_coeff_save)}")
        
         